from pathlib import Path

from ax.api.client import Client
from src.metric import ax_metrics_from_file
from src.runner import after_optimization, config_client, run_experiment

# 1. 初始化客户端
client = Client()

preexisting_trials_dir = "data/preexisting_trials"
new_data_dir = "data"
results_dir = "data/ax_snapshot"

# 2. 配置搜索空间, 优化目标, 当前转速
tail_temperature_range = [25, 45, 55, 65, 75]
motor_rpm = 30
optimize_metric = "efficiency"

# 配置优化器, 如果不加载已有实验，则将 preexisting_trials_dir 设为 None
config_client(
    client=client,
    tail_temperature_range=tail_temperature_range,
    motor_rpm=motor_rpm,
    optimize_metric=optimize_metric,
    preexisting_trials_dir=preexisting_trials_dir,
)

# 3. 进行 20 次试验的实验：从 Ax 获取每个试验，评估目标函数，并将数据记录回 Ax
for i in range(20):
    # 每次最多运行几个优化max_trials
    trials = client.get_next_trials(max_trials=1)
    trial_index, params = next(iter(trials.items()))

    # 运行实验并获取数据文件路径
    data_file_path = run_experiment(i, params, motor_rpm, new_data_dir)

    # 如果用户跳过了实验，则跳过此次迭代
    if data_file_path == "skip":
        print("⚠️  实验被跳过，将不记录此次试验结果")
        client.mark_trial_failed(trial_index=trial_index)
        continue
    elif data_file_path == "exit":
        print("⚠️  实验被用户退出，终止优化过程")
        client.mark_trial_failed(trial_index=trial_index)
        break
    elif isinstance(data_file_path, Path):
        # 从数据文件中提取指标
        metric_dict = ax_metrics_from_file(
            data_file_path, optimize_metric, motor_rpm
        )

    # 将数据记录回 Ax
    client.complete_trial(trial_index=trial_index, raw_data=metric_dict)

# 4. 优化结束后处理
after_optimization(
    client=client,
    results_dir=results_dir,
    optimize_metric=optimize_metric,
    motor_rpm=motor_rpm,
)
